package fasta

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"
	"strconv"
	"cs466/findmotif/motif"
)

// read the FASTA sequence, stripping newlines and the leading comment
func ReadFasta(r io.Reader) (seqs [][]byte, names []string, err error) {
	seqs = make([][]byte, 0)
	names = make([]string, 0)
	var current_seq []byte

	b := bufio.NewReader(r)

	var line []byte
	for line, err = b.ReadBytes('\n'); err == nil; line, err = b.ReadBytes('\n') {
		switch line[0] {
		case '>':
			names = append(names, string(line[1:]))
			if current_seq != nil {
				seqs = append(seqs, current_seq)
			}
			current_seq = make([]byte, 0)
		default:
			current_seq = append(current_seq, line[:len(line)-1]...)
		}
	}
	if err != io.EOF {
		return
	}
	seqs = append(seqs, current_seq)
	err = nil
	return
}

// load the FASTA sequence from the specified file
func LoadFasta(fname string) ([][]byte, []string, error) {
	f, err := os.Open(fname)
	if err != nil {
		return nil, nil, err
	}
	defer f.Close()
	return ReadFasta(f)
}

func ReadMotif(r io.Reader) (pwm [][]float64, err error) {
	pwm = make([][]float64, 0)
	var current_col []float64

	b := bufio.NewReader(r)

	// skip first line
	_,err = b.ReadString('\n')

	var line string
	for line, err = b.ReadString('\n'); err == nil; line, err = b.ReadString('\n') {
		col := strings.Split(strings.TrimSpace(line), "\t")
		current_col = make([]float64, len(col))
		for i, v := range col {
			current_col[i], err = strconv.ParseFloat(v, 64)
		}
		pwm = append(pwm, current_col)
	}
	if err != io.EOF {
		return
	}
	pwm = append(pwm, current_col)
	err = nil

	return
}

func LoadMotif(fname string) ([][]float64, error) {
	f, err := os.Open(fname)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return ReadMotif(f)
}

func LoadMotifLength(fname string) (uint, error) {
	var ml uint
	f, err := os.Open(fname)
	if err != nil {
		return 0, err
	}
	defer f.Close()

	_, err = fmt.Fscanf(f, "%d", &ml)
	if err != nil {
		return 0, err
	}

	return ml, nil
}

func LoadSites(fname string) (sites []int, err error) {
	f, err := os.Open(fname)
	if err != nil {
		return
	}
	defer f.Close()

	b := bufio.NewReader(f)

	line, err := b.ReadString('\n')
	if err != io.EOF {
		return
	}

	line = line[1:len(line)-1]

	toks := strings.Split(line, " ")

	sites = make([]int, len(toks))

	for i, v := range toks {
		sites[i], err = strconv.Atoi(v)
		if err != nil {
			return
		}
	}

	return
}

// dump the FASTA sequence to the provided writer, limiting line lengths
func WriteFasta(w io.Writer, comment string, sequence []byte) (err error) {
	_, err = fmt.Fprintf(w, ">%s\n", comment)
	if err != nil {
		return err
	}

	for len(sequence) >= 80 {
		_, err = fmt.Fprintf(w, "%s\n", string(sequence[:80]))
		if err != nil {
			return err
		}

		sequence = sequence[80:]
	}

	_, err = fmt.Fprintf(w, "%s\n", string(sequence))

	return
}

// dump the whole map of FASTA sequences to the provided writer
func WriteFastaMap(w io.Writer, seqs map[string][]byte) (err error) {
	for comment, seq := range seqs {
		err = WriteFasta(w, comment, seq)
		if err != nil {
			return
		}
	}

	return nil
}

func WriteFastaSlice(w io.Writer, seqs [][]byte) (err error) {
	for i, seq := range seqs {
		err = WriteFasta(w, fmt.Sprint(i), seq)
		if err != nil {
			return
		}
	}

	return nil
}

func WriteMotif(w io.Writer, pwm motif.PWM, ml uint, sc float64) (err error) {
	_, err = fmt.Fprintf(w, ">MOTIF1\t%d\n", ml)
	if err != nil {
		return err
	}
	for _, column := range pwm {
		for i, prob := range column {
			_, err = fmt.Fprint(w, sc * prob)
			if err != nil {
				return err
			}
			if i < len(column) - 1 {
				_, err = fmt.Fprint(w, "\t")
				if err != nil {
					return err
				}
			}
		}
		_, err = fmt.Fprint(w, "\n")
		if err != nil {
			return err
		}
	}

	return nil
}
